# MPC单输入，两个状态变量示例。相当于LQR调节器，xd=0。无约束优化。
from datetime import datetime
import numpy as np
import math
import matplotlib.pyplot as plt
import pandas as pd
import control as ct
import time
import scipy.linalg as la
from scipy.signal import StateSpace
from scipy.optimize import minimize
import quadprog

# 定义MPC转化函数
def MPC_Matrices(A,B,Q,R,F,N) :
    # 获取行数
    n = len(A)
    # 获取列数
    p =  B.shape[1]
    M = np.vstack((np.eye(n), np.zeros((N*n,n))))
    C = np.zeros(((N+1)*n,N*p))
    tmp = np.eye(n)
    
    for i in range(1,N+1):
        # 生成从 i*n 到 (i+1)*n-1 的行索引
        rows = np.arange(i*n, (i+1)*n)
        # print((tmp@B).shape,rows)
        C[rows,:] = np.hstack((tmp@B,C[rows-n, :-p]))
        tmp = A@tmp
        M[rows,:] = tmp
    
    Q_bar = np.kron(np.eye(N), Q)#张量积
    # Q_bar与F的块对角阵
    Q_bar = np.block([[Q_bar, np.zeros((N*Q.shape[0], F.shape[1]))],
                  [np.zeros((F.shape[0], N*Q.shape[1])), F]])
    R_bar = np.kron(np.eye(N), R)
    E = C.T@Q_bar@M
    H = C.T@Q_bar@C+R_bar
    return E,H

# 二次规划

def Prediction(x_k, E, H, N, p):
    """
    使用二次规划求解控制输入 u_k

    参数:
        x_k: 当前状态向量 (n x 1)
        E: 矩阵，用于计算目标函数 (m x n)
        H: 二次规划的正定矩阵 (m x m)
        N: 预测时域长度
        p: 控制输入的维度

    返回:
        u_k: 当前时刻的控制输入 (p x 1)
    """
    # 计算目标函数的一次项系数
    q = E @ x_k  # q = E * x_k
    U_k = np.zeros((N*p,1))
    # 调用 quadprog 求解二次规划
    U_k = quadprog.solve_qp(H, q)#元组中的第一个元素为最优化的输入的np数组

    # 提取元组第一个控制输入
    
    u_k = U_k[0][0]
    return -u_k

# 1.定义系统
A = np.array([[1,0.1],[0,2]])
B = np.array([[0],[1/2]])
print("当前时间是:",datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
# 获取行数
n = len(A)
# 获取列数
p =  B.shape[1]
# 设置权重
Q = np.array([[1,0],[0,1]])
F = np.array([[1,0],[0,1]])
R= np.array([[0.1]])
k_steps=100
# 2.初始化系统
x = np.array([[20],[-20]]).reshape(1,-1)
X_K = np.zeros((n,k_steps+1))

X_K[:,0] = x
# 输入
U_K = np.zeros((p,k_steps))

N = 5
E,H = MPC_Matrices(A,B,Q,R,F,N)


for k in range(1,k_steps+1):
    U_K[:,k-1] = Prediction(X_K[:,k-1],E,H,N,p).reshape(p,1)
    X_K[:,k] = (A@(X_K[:,k-1]).reshape(-1,1)+(B@(U_K[:,k-1])).reshape(-1,1)).flatten()


# # 4.制图
plt.figure()
plt.subplot(2,1,1)
plt.title('Input and State Plot')
plt.xlabel('steps')
plt.ylabel('State')
plt.plot(X_K[0,:],label=f'State 1')
plt.plot(X_K[1,:],label=f'State 2')

plt.legend()
plt.subplot(2,1,2)
plt.title('u Exception Plot')
plt.xlabel('steps')
plt.ylabel('u')
plt.plot(U_K[0,:],label=f'u 1')
plt.legend()
    
plt.legend()
plt.grid()
plt.show()